import os
import random
import argparse
from shutil import copy

def make_parser():
    parser = argparse.ArgumentParser("onnx deploy")
    parser.add_argument("--txt_path", type=str, default="./BigData/message.txt", help="txt path")
    parser.add_argument("--tag_path", type=str, default="./BigData/tag.txt")
    parser.add_argument("--input_path", type=str, default="./BigData", help="input img path")
    parser.add_argument("--output_path", type=str, default="./BigData/dataset")
    return parser

def main():
    args = make_parser().parse_args()
    txt_path = args.txt_path
    tag_path = args.tag_path
    input_path = args.input_path
    output_path = args.output_path

    input_o_path = input_path + "/image_original"
    input_m_path = input_path + "/image_message"

    if not os.path.exists(output_path):
        os.mkdir(output_path)
    
    img_dir = os.path.join(output_path, "tarin")
    output_txt = os.path.join(output_path, "train.txt")
    tag_file = open(tag_path, 'r')
    tags = tag_file.readline()
    tags = tags.rstrip("\n").split()

    if not os.path.exists(img_dir):
        os.mkdir(img_dir)
    
    img_o_dir = img_dir + "/image_original"
    img_m_dir = img_dir + "/image_message"
    if not os.path.exists(img_o_dir):
        os.mkdir(img_o_dir)
    if not os.path.exists(img_m_dir):
        os.mkdir(img_m_dir)

    output_txt_file = open(output_txt, 'w')

    txt_file = open(txt_path, 'r')
    lines = txt_file.readlines()

    total_cnt = len(lines)
    cnt = 0
    train_num = int(0.8 * total_cnt)
    random.shuffle(lines)

    for line in lines:
        cnt += 1
        if cnt == train_num:
            img_dir = os.path.join(output_path, "test")
            if not os.path.exists(img_dir):
                os.mkdir(img_dir)
            img_o_dir = img_dir + "/image_original"
            img_m_dir = img_dir + "/image_message"
            if not os.path.exists(img_o_dir):
                os.mkdir(img_o_dir)
            if not os.path.exists(img_m_dir):
                os.mkdir(img_m_dir)
            output_txt = os.path.join(output_path, "test.txt")
            output_txt_file.close()
            output_txt_file = open(output_txt, 'w')

        img_name = line[2:9]
        label = line[11]
        for i in range(len(tags)):
            if label == tags[i]:
                label = i
                break
        
        img_name = img_name + ".jpg"
        #img_path = os.path.join(input_path, img_name)
        img_o_path = os.path.join(input_o_path, img_name)
        img_m_path = os.path.join(input_m_path, img_name)

        output_txt_file.write(img_o_dir + "/" + img_name + " " + img_m_dir + "/" + img_name + " " + str(label) + "\n")
        copy(img_o_path, img_o_dir)
        copy(img_m_path, img_m_dir)
        print(img_name, " handle")
    print("target total cnt: ", total_cnt, "    total cnt: ", cnt)

if __name__ == "__main__":
    main()